"""MongoDB + Redis 存储实现。"""

from __future__ import annotations

import json
from datetime import datetime, timedelta, timezone
from typing import List, Optional

from bson import ObjectId
from bson.errors import InvalidId
from pymongo import ASCENDING, MongoClient, ReturnDocument
from pymongo.errors import ConfigurationError
from redis import Redis

from app.core.config import settings
from app.core.security import decrypt_secret, encrypt_secret
from app.schemas.certificates import (
    Certificate,
    CertificateAutoApplyRequest,
    CertificateCreate,
    CertificateRenewRequest,
    CertificateStatus,
    CertificateUploadRequest,
)
from app.schemas.credentials import (
    Credential,
    CredentialCreate,
    CredentialDescribeRequest,
    CredentialListRequest,
    CredentialUpdate,
)
from app.schemas.dns import (
    DNSRecord,
    DNSRecordBatchDeleteRequest,
    DNSRecordCreate,
    DNSRecordIdentifier,
    DNSRecordListRequest,
    DNSRecordStatus,
    DNSRecordType,
    DNSRecordUpdate,
    DomainListRequest,
    DomainSummary,
)


def _now() -> datetime:
    return datetime.now(timezone.utc)


def _mask_secret(secret: str) -> str:
    if len(secret) <= 4:
        return "*" * len(secret)
    return secret[:2] + "****" + secret[-2:]


def _object_id(value: str) -> ObjectId | str:
    try:
        return ObjectId(value)
    except (InvalidId, TypeError):
        # If not a valid ObjectId, assume it's a custom ID (e.g. Cloud ID)
        return value


class MongoStore:
    def __init__(self) -> None:
        self.client = MongoClient(settings.database_url)
        self.db = self._resolve_database()
        self.credentials = self.db["credentials"]
        self.dns_records = self.db["dns_records"]
        self.certificates = self.db["certificates"]
        self.acme_accounts = self.db["acme_accounts"]
        self.redis = Redis.from_url(settings.redis_url, decode_responses=True)
        self._ensure_indexes()

    @staticmethod
    def _object_id(value: str) -> ObjectId | str:
        try:
            return ObjectId(value)
        except (InvalidId, TypeError):
            # If not a valid ObjectId, assume it's a custom ID (e.g. Cloud ID)
            return value

    def _resolve_database(self):
        try:
            database = self.client.get_default_database()
            if database is not None:
                return database
        except ConfigurationError:
            pass
        db_name = settings.database_url.rsplit("/", 1)[-1].split("?")[0] or "cert_dns_hub"
        return self.client.get_database(db_name)

    def _ensure_indexes(self) -> None:
        self.credentials.create_index([("provider", ASCENDING)])
        self.credentials.create_index([("account_name", ASCENDING)])
        self.dns_records.create_index([("provider", ASCENDING), ("domain", ASCENDING)])
        self.dns_records.create_index([("domain", ASCENDING), ("name", ASCENDING)])
        self.certificates.create_index([("provider", ASCENDING)])
        self.certificates.create_index([("expires_at", ASCENDING)])
        self.acme_accounts.create_index([("email", ASCENDING)], unique=True)

    def _domain_cache_key(self, provider: Optional[str]) -> str:
        return f"dns:domains:{provider or 'all'}"

    def _invalidate_domain_cache(self) -> None:
        keys = list(self.redis.scan_iter(match="dns:domains:*"))
        if keys:
            self.redis.delete(*keys)

    # ----- ACME Account Helpers -----
    def get_default_acme_account(self) -> dict:
        """获取默认 ACME 账户，如果不存在则创建。"""
        # For simplicity, we use a fixed email or identifier for the default account
        email = "default@example.com"
        account = self.acme_accounts.find_one({"email": email})
        if account:
            return account
        
        # Generate new key pair (RSA 2048)
        from cryptography.hazmat.primitives import serialization
        from cryptography.hazmat.primitives.asymmetric import rsa
        
        key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
        private_key_pem = key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=serialization.NoEncryption(),
        ).decode("utf-8")
        public_key_pem = key.public_key().public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo,
        ).decode("utf-8")
        
        account_data = {
            "email": email,
            "private_key": private_key_pem,
            "public_key": public_key_pem,
            "created_at": _now(),
            "registration_uri": None, # To be filled after registration
        }
        self.acme_accounts.insert_one(account_data)
        return account_data

    def update_acme_account_registration(self, email: str, uri: str) -> None:
        self.acme_accounts.update_one({"email": email}, {"$set": {"registration_uri": uri}})

    # ----- Credential Helpers -----
    @staticmethod
    def _to_credential(doc: dict) -> Credential:
        return Credential(
            id=str(doc["_id"]),
            provider=doc["provider"],
            account_name=doc["account_name"],
            access_key=doc["access_key"],
            description=doc.get("description"),
            secret_preview=doc["secret_preview"],
        )

    # ... (Credentials and DNS methods remain unchanged) ...
    # I will skip repeating them to save tokens, but I need to be careful with replace_file_content.
    # Since I'm replacing a large chunk, I should probably use multi_replace or be precise.
    # The user asked to update datastore.py.
    # I will use replace_file_content to update __init__ and add ACME methods, 
    # and then update certificate methods.
    
    # Wait, I can't skip code in replace_file_content.
    # I will use multi_replace_file_content to target specific blocks.

    # Let's abort this tool call and use multi_replace_file_content.

    def _resolve_database(self):
        try:
            database = self.client.get_default_database()
            if database is not None:
                return database
        except ConfigurationError:
            pass
        db_name = settings.database_url.rsplit("/", 1)[-1].split("?")[0] or "cert_dns_hub"
        return self.client.get_database(db_name)

    def _ensure_indexes(self) -> None:
        self.credentials.create_index([("provider", ASCENDING)])
        self.credentials.create_index([("account_name", ASCENDING)])
        self.dns_records.create_index([("provider", ASCENDING), ("domain", ASCENDING)])
        self.dns_records.create_index([("domain", ASCENDING), ("name", ASCENDING)])
        self.certificates.create_index([("provider", ASCENDING)])
        self.certificates.create_index([("expires_at", ASCENDING)])

    def _domain_cache_key(self, provider: Optional[str]) -> str:
        return f"dns:domains:{provider or 'all'}"

    def _invalidate_domain_cache(self) -> None:
        keys = list(self.redis.scan_iter(match="dns:domains:*"))
        if keys:
            self.redis.delete(*keys)

    # ----- Credential Helpers -----
    @staticmethod
    def _to_credential(doc: dict) -> Credential:
        return Credential(
            id=str(doc["_id"]),
            provider=doc["provider"],
            account_name=doc["account_name"],
            access_key=doc["access_key"],
            description=doc.get("description"),
            secret_preview=doc["secret_preview"],
        )

    # ----- Credentials -----
    def create_credential(self, payload: CredentialCreate) -> Credential:
        data = payload.model_dump()
        secret_key = data.pop("secret_key")
        data["encrypted_secret"] = encrypt_secret(secret_key)
        data["secret_preview"] = _mask_secret(secret_key)
        result = self.credentials.insert_one(data)
        doc = self.credentials.find_one({"_id": result.inserted_id})
        return self._to_credential(doc)

    def update_credential(self, payload: CredentialUpdate) -> Credential:
        oid = _object_id(payload.id)
        update_data = payload.model_dump(exclude_none=True)
        update_data.pop("id", None)
        if "secret_key" in update_data:
            raw_secret = update_data.pop("secret_key")
            update_data["encrypted_secret"] = encrypt_secret(raw_secret)
            update_data["secret_preview"] = _mask_secret(raw_secret)
        updated = self.credentials.find_one_and_update(
            {"_id": oid},
            {"$set": update_data},
            return_document=ReturnDocument.AFTER,
        )
        if not updated:
            raise KeyError("credential not found")
        return self._to_credential(updated)

    def delete_credential(self, credential_id: str) -> None:
        oid = _object_id(credential_id)
        result = self.credentials.delete_one({"_id": oid})
        if result.deleted_count == 0:
            raise KeyError("credential not found")

    def list_credentials(self, payload: Optional[CredentialListRequest]) -> List[Credential]:
        query = {}
        if payload and payload.provider:
            query["provider"] = payload.provider
        docs = self.credentials.find(query)
        return [self._to_credential(doc) for doc in docs]

    def describe_credential(self, payload: CredentialDescribeRequest) -> Credential:
        oid = _object_id(payload.id)
        doc = self.credentials.find_one({"_id": oid})
        if not doc:
            raise KeyError("credential not found")
        return self._to_credential(doc)

    def get_decrypted_credential(self, credential_id: str) -> tuple[str, str, str]:
        """获取解密后的凭据信息 (provider, access_key, secret_key)。"""
        oid = _object_id(credential_id)
        doc = self.credentials.find_one({"_id": oid})
        if not doc:
            raise KeyError("credential not found")
        
        secret_key = decrypt_secret(doc["encrypted_secret"])
        return doc["provider"], doc["access_key"], secret_key

    # ----- DNS Helpers -----
    @staticmethod
    def _to_dns_record(doc: dict) -> DNSRecord:
        status_value = doc.get("status", DNSRecordStatus.active.value)
        return DNSRecord(
            id=str(doc["_id"]),
            provider=doc["provider"],
            domain=doc["domain"],
            record_type=DNSRecordType(doc["record_type"]),
            name=doc["name"],
            value=doc["value"],
            ttl=doc["ttl"],
            status=DNSRecordStatus(status_value),
        )

    # ----- DNS -----
    def get_dns_record(self, record_id: str) -> DNSRecord:
        oid = _object_id(record_id)
        doc = self.dns_records.find_one({"_id": oid})
        if not doc:
            raise KeyError("dns record not found")
        return self._to_dns_record(doc)

    def create_dns_record(self, payload: DNSRecordCreate, extra_data: Optional[dict] = None) -> DNSRecord:
        data = payload.model_dump()
        if extra_data:
            data.update(extra_data)
        
        # If extra_data has 'id', we might want to use it as _id if it's a valid ObjectId, 
        # or just store it as 'cloud_id'.
        # For this demo, if 'id' is passed in extra_data, we try to use it as _id if possible,
        # but Cloud IDs are usually strings (integers or UUIDs).
        # MongoDB _id can be any type.
        # Let's use the Cloud ID as _id if provided.
        if extra_data and "id" in extra_data:
            data["_id"] = extra_data["id"]
            
        try:
            result = self.dns_records.insert_one(data)
            doc = self.dns_records.find_one({"_id": result.inserted_id})
        except Exception:
            # Fallback if duplicate key or other error, or just let it raise
            raise
            
        self._invalidate_domain_cache()
        return self._to_dns_record(doc)

    def update_dns_record(self, payload: DNSRecordUpdate) -> DNSRecord:
        oid = _object_id(payload.id)
        updates = payload.model_dump(exclude_none=True)
        updates.pop("id", None)
        updated = self.dns_records.find_one_and_update(
            {"_id": oid},
            {"$set": updates},
            return_document=ReturnDocument.AFTER,
        )
        if not updated:
            raise KeyError("dns record not found")
        self._invalidate_domain_cache()
        return self._to_dns_record(updated)

    def delete_dns_record(self, payload: DNSRecordIdentifier) -> None:
        oid = _object_id(payload.id)
        result = self.dns_records.delete_one({"_id": oid})
        if result.deleted_count == 0:
            raise KeyError("dns record not found")
        self._invalidate_domain_cache()

    def batch_delete_dns_record(self, payload: DNSRecordBatchDeleteRequest) -> int:
        object_ids = []
        for record_id in payload.ids:
            try:
                object_ids.append(ObjectId(record_id))
            except InvalidId:
                continue
        if not object_ids:
            return 0
        result = self.dns_records.delete_many({"_id": {"$in": object_ids}})
        if result.deleted_count:
            self._invalidate_domain_cache()
        return result.deleted_count

    def list_dns_records(self, payload: Optional[DNSRecordListRequest]) -> List[DNSRecord]:
        query = {}
        if payload:
            if payload.domain:
                query["domain"] = payload.domain
            if payload.provider:
                query["provider"] = payload.provider
        docs = self.dns_records.find(query)
        return [self._to_dns_record(doc) for doc in docs]

    def list_domains(self, payload: Optional[DomainListRequest]) -> List[DomainSummary]:
        # provider = payload.provider if payload else None
        # Since provider is removed from request, we list all or rely on credential_id (handled in endpoint)
        # For local listing, we just list all for now or use cache key 'all'
        provider = None
        cache_key = self._domain_cache_key(provider)
        cached = self.redis.get(cache_key)
        if cached:
            data = json.loads(cached)
            return [DomainSummary(**item) for item in data]

        pipeline = []
        if provider:
            pipeline.append({"$match": {"provider": provider}})
        pipeline.extend(
            [
                {
                    "$group": {
                        "_id": {"domain": "$domain", "provider": "$provider"},
                        "record_count": {"$sum": 1},
                    }
                },
                {
                    "$project": {
                        "_id": 0,
                        "domain": "$_id.domain",
                        "provider": "$_id.provider",
                        "record_count": "$record_count",
                    }
                },
            ]
        )
        docs = list(self.dns_records.aggregate(pipeline))
        self.redis.setex(cache_key, 30, json.dumps(docs))
        return [DomainSummary(**doc) for doc in docs]

    def set_dns_record_status(self, record_id: str, status: DNSRecordStatus) -> DNSRecord:
        oid = _object_id(record_id)
        updated = self.dns_records.find_one_and_update(
            {"_id": oid},
            {"$set": {"status": status.value}},
            return_document=ReturnDocument.AFTER,
        )
        if not updated:
            raise KeyError("dns record not found")
        self._invalidate_domain_cache()
        return self._to_dns_record(updated)

    # ----- Certificate Helpers -----
    @staticmethod
    def _to_certificate(doc: dict) -> Certificate:
        return Certificate(
            id=str(doc["_id"]),
            provider=doc.get("provider"),
            credential_id=doc.get("credential_id"),
            common_name=doc["common_name"],
            domains=doc.get("domains", []),
            status=CertificateStatus(doc["status"]),
            auto_renew=doc.get("auto_renew", True),
            order_url=doc.get("order_url"),
            created_at=doc["created_at"],
            updated_at=doc["updated_at"],
            expires_at=doc.get("expires_at"),
        )

    # ----- Certificates -----
    def create_certificate(self, payload: CertificateCreate) -> Certificate:
        now = _now()
        data = payload.model_dump()
        data.update(
            status=CertificateStatus.pending_validation.value,
            created_at=now,
            updated_at=now,
            expires_at=None,
        )
        result = self.certificates.insert_one(data)
        doc = self.certificates.find_one({"_id": result.inserted_id})
        return self._to_certificate(doc)

    def auto_apply_certificate(self, payload: CertificateAutoApplyRequest) -> Certificate:
        # We don't have provider in payload anymore, we have credential_id
        # But we need provider to store in certificate? 
        # Actually schema says provider is optional.
        # We can fetch provider from credential if needed, or just store credential_id.
        # Let's store credential_id.
        
        domains = payload.domains or []
        common_name = domains[0] if domains else "unknown"
        
        create_payload = CertificateCreate(
            credential_id=payload.credential_id,
            common_name=common_name,
            domains=domains,
            auto_renew=True,
        )
        certificate = self.create_certificate(create_payload)
        return certificate

    def renew_certificate(self, payload: CertificateRenewRequest) -> Certificate:
        oid = _object_id(payload.id)
        updated = self.certificates.find_one_and_update(
            {"_id": oid},
            {
                "$set": {
                    "updated_at": _now(),
                    # Status might change to pending if we re-verify?
                    # For now just update timestamp, actual logic handled in service/endpoint
                }
            },
            return_document=ReturnDocument.AFTER,
        )
        if not updated:
            raise KeyError("certificate not found")
        return self._to_certificate(updated)

    def upload_certificate(self, payload: CertificateUploadRequest) -> Certificate:
        now = _now()
        data = dict(
            provider=payload.provider,
            common_name=payload.common_name,
            domains=[payload.common_name],
            status=CertificateStatus.issued.value,
            auto_renew=False,
            created_at=now,
            updated_at=now,
            expires_at=now + timedelta(days=365), # Estimate
            certificate_pem=payload.certificate_pem,
            private_key_pem=payload.private_key_pem,
        )
        result = self.certificates.insert_one(data)
        doc = self.certificates.find_one({"_id": result.inserted_id})
        return self._to_certificate(doc)

    def list_certificates(self) -> List[Certificate]:
        docs = self.certificates.find()
        return [self._to_certificate(doc) for doc in docs]

    def describe_certificate(self, cert_id: str) -> Certificate:
        oid = _object_id(cert_id)
        doc = self.certificates.find_one({"_id": oid})
        if not doc:
            raise KeyError("certificate not found")
        return self._to_certificate(doc)

    def delete_certificate(self, cert_id: str) -> None:
        oid = _object_id(cert_id)
        result = self.certificates.delete_one({"_id": oid})
        if result.deleted_count == 0:
            raise KeyError("certificate not found")

    def update_certificate_status(self, cert_id: str, status: CertificateStatus, 
                                  cert_pem: Optional[str] = None, 
                                  key_pem: Optional[str] = None,
                                  expires_at: Optional[datetime] = None) -> Certificate:
        oid = _object_id(cert_id)
        update_data = {
            "status": status.value,
            "updated_at": _now(),
        }
        if cert_pem:
            update_data["certificate_pem"] = cert_pem
        if key_pem:
            update_data["private_key_pem"] = key_pem
        if expires_at:
            update_data["expires_at"] = expires_at
            
        updated = self.certificates.find_one_and_update(
            {"_id": oid},
            {"$set": update_data},
            return_document=ReturnDocument.AFTER,
        )
        if not updated:
            raise KeyError("certificate not found")
        return self._to_certificate(updated)

    def _update_certificate_status(self, cert_id: str, status: CertificateStatus) -> Certificate:
        # Deprecated internal helper, redirected to new one
        return self.update_certificate_status(cert_id, status)


store = MongoStore()
